from read_CloudSat import reader
import numpy as np
import numpy.ma as ma
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
import cartopy.crs as ccrs
from cartopy.mpl.ticker import LongitudeFormatter, LatitudeFormatter
from datetime import datetime
import pandas as pd
import netCDF4 as nc
from scipy.spatial import cKDTree


def set_map(ax):
    '''设置地图.'''
    proj = ccrs.PlateCarree()
    ax.coastlines(lw=0.5)
    xticks = np.arange(-180, 180 + 60, 60)
    yticks = np.arange(-90, 90 + 30, 30)
    ax.set_xticks(xticks, crs=proj)
    ax.set_yticks(yticks, crs=proj)
    ax.xaxis.set_major_formatter(LongitudeFormatter())
    ax.yaxis.set_major_formatter(LatitudeFormatter())
    ax.tick_params('both', labelsize='x-small')
    ax.set_global()


def draw_track(ax, lon1D, lat1D):
    '''根据经纬度画出轨迹,并标识出起点.'''
    ax.plot(lon1D, lat1D, lw=2, color='b', transform=ccrs.Geodetic())
    ax.plot(lon1D[0], lat1D[0], 'ro', ms=3, transform=ccrs.PlateCarree())
    ax.text(
        lon1D[0] + 5, lat1D[0], 'start', color='r', fontsize='x-small',
        transform=ccrs.PlateCarree()
    )


def draw_cross_section(ax, lon, hgt, data, cth_sorted, cbh_sorted, latitude_sorted):
    '''画出data的经度-高度剖面，并添加风云数据的云底和云顶高度折线图.'''
    im = ax.pcolormesh(lon, hgt, data, cmap='jet', shading='nearest')
    ax.set_ylim(0, 20)
    ax.set_xlim(lon.min(), lon.max())
    ax.tick_params(labelsize='x-small')
    ax.set_xlabel('Longitude', fontsize='x-small')
    ax.set_ylabel('Height [km]', fontsize='x-small')
    divider = make_axes_locatable(ax)
    cax = divider.append_axes('bottom', size='10%', pad=0.42)
    cbar = plt.colorbar(im, cax=cax, extend='both', orientation='horizontal')
    cbar.ax.tick_params(labelsize='x-small')
    cbar.set_label('dBZe', fontsize='x-small')
    # 使用df_sorted中的经度作为横坐标
    ax.plot(df_sorted['Longitude'], cth_sorted, color='#0047AB', label='CTH vs Longitude',linewidth=0.5)
    ax.plot(df_sorted['Longitude'], cbh_sorted, color='#FF8C00', label='CBH vs Longitude',linewidth=0.5)
    ax.legend(fontsize='x-small')
    # # 绘制风云数据的云底和云顶高度折线图
    # ax.plot(latitude_sorted, cth_sorted, color='blue', label='CTH vs Latitude')
    # ax.plot(latitude_sorted, cbh_sorted, color='red', label='CBH vs Latitude')
    # ax.legend(fontsize='x-small')


def draw_elevation(ax, lon, elv):
    '''画出地形抬升.'''
    ax.fill_between(lon, elv, color='gray')


if __name__ == '__main__':
    # 读取CloudSat数据
    fname = (
        '/mnt/datastore/liudddata/cloudsat_data/cloudsat_GEOPROF/cloudsat_2020_1_4/2020112031512_74484_CS_2B-GEOPROF_GRANULE_P1_R05_E09_F00.hdf')
    f = reader(fname)
    lon, lat, elv = f.read_geo()
    data = f.read_sds('Radar_Reflectivity')
    height = f.read_sds('Height')
    time = f.read_time(datetime=True)
    f.close()

    # 手动指定开始时间和结束时间
    start_datetime = datetime(2020, 4, 21, 4, 0, 0)
    end_datetime = datetime(2020, 4, 21, 4, 15, 0)
    # 寻找时间范围内的索引
    time_series = pd.to_datetime(time)
    start_index, end_index = np.searchsorted(time_series, [start_datetime, end_datetime])

    # 选择时间范围内的数据
    rad_refl_t = np.transpose(data[start_index:end_index])
    rad_refl_t[rad_refl_t <= -20] = np.nan  # 过滤数据

    # 获取对应时间范围内的高度数据，单位转换为km
    hgt = np.mean(height[start_index:end_index], axis=0) / 1000
    lat_t = lat[start_index:end_index]
    lon_t = lon[start_index:end_index]
    lat_t[np.isnan(lat_t)] = np.nan
    lon_t[np.isnan(lon_t)] = np.nan

    # 加载风云数据
    filename = "/mnt/datastore/liudddata/result/20200104new/2020042104_predicted_2d.nc"
    cth_file = nc.Dataset(filename, 'r')
    cth_data = cth_file.variables['cth'][:]
    cbh_data = cth_file.variables['predicted'][:]
    latitudes = []
    longitudes = []
    cth_values = []
    cbh_values = []

    # 加载地理坐标数据
    coord_file_name = 'FY4A_coordinates.nc'
    coord_file_open = nc.Dataset(coord_file_name, 'r')
    lat_fy = coord_file_open.variables['lat'][:, :].T
    lon_fy = coord_file_open.variables['lon'][:, :].T
    lat_fy[np.isnan(lat_fy)] = -90.
    lon_fy[np.isnan(lon_fy)] = 360.
    coord_file_open.close()

    # 将 lat_fy 和 lon_fy 转换为笛卡尔坐标系以便使用 cKDTree
    fy_coords = np.deg2rad(np.vstack([lat_fy.ravel(), lon_fy.ravel()]).T)
    fy_coords_cartesian = np.array([np.cos(fy_coords[:, 0]) * np.cos(fy_coords[:, 1]),
                                    np.cos(fy_coords[:, 0]) * np.sin(fy_coords[:, 1]),
                                    np.sin(fy_coords[:, 0])]).T

    # 确保 fy_coords_cartesian 中没有 NaN 或无穷大/无穷小值
    is_finite = np.isfinite(fy_coords_cartesian).all(axis=1)

    # 仅保留所有值都是有限的行
    fy_coords_cartesian_clean = fy_coords_cartesian[is_finite]

    # 现在可以安全地创建 cKDTree
    tree = cKDTree(fy_coords_cartesian_clean)

    # 对于每个 lat_t 和 lon_t，找到2km范围内的 lat_fy 和 lon_fy
    radius = 2 / 6371  # 2km 转换为弧度
    for lat, lon in zip(lat_t, lon_t):
        point = np.deg2rad(np.array([lat, lon]))
        point_cartesian = np.array([np.cos(point[0]) * np.cos(point[1]),
                                    np.cos(point[0]) * np.sin(point[1]),
                                    np.sin(point[0])])

        # 执行半径搜索
        indices = tree.query_ball_point(point_cartesian, r=radius)

        # 处理找到的点
        if indices:
            for idx in indices:
                lat_idx, lon_idx = np.unravel_index(idx, lat_fy.shape)
                latitudes.append(lat_fy[lat_idx, lon_idx])
                longitudes.append(lon_fy[lat_idx, lon_idx])
                cth_values.append(cth_data[lat_idx, lon_idx] * 0.001)
                cbh_values.append(cbh_data[lat_idx, lon_idx] * 0.001)
        else:
            # 如果没有找到匹配，添加NaN作为占位符
            latitudes.append(np.nan)
            longitudes.append(np.nan)
            cth_values.append(np.nan)
            cbh_values.append(np.nan)

    fy_data = {
        'Latitude': latitudes,
        'Longitude': longitudes,
        'CTH': cth_values,
        'CBH': cbh_values
    }

    df = pd.DataFrame(fy_data)

    # 过滤掉包含NaN值的行
    df.dropna(inplace=True)
    # 按照纬度排序
    df_sorted = df.sort_values(by='Latitude')

    # 获取排序后的纬度和CTH
    latitude_sorted = df_sorted['Latitude']
    cth_sorted = df_sorted['CTH']
    cbh_sorted = df_sorted['CBH']

    # 高度单位全部转为km.
    elv /= 1000

    fig = plt.figure(dpi=200, figsize=(6, 6))
    ax1 = fig.add_axes([0.25, 0.45, 0.5, 0.5], projection=ccrs.PlateCarree())
    ax2 = fig.add_axes([0.25, 0.2, 0.5, 0.3])

    set_map(ax1)
    draw_track(ax1, lon_t, lat_t)

    # 把granule的起止时间标上.
    start_str = start_datetime.strftime('%Y-%m-%d %H:%M')
    end_str = end_datetime.strftime('%Y-%m-%d %H:%M')
    ax1.set_title(f'From {start_str} to {end_str}', fontsize='small')

    # draw_cross_section(ax2, lon_t, hgt, rad_refl_t, cth_sorted, cbh_sorted, latitude_sorted)
    draw_cross_section(ax2, lon_t, hgt, rad_refl_t, cth_sorted, cbh_sorted, df_sorted)
    draw_elevation(ax2, lon_t, elv[start_index:end_index])
    ax2.set_title('Radar Reflectivity Factor', fontsize='small')

    fig.savefig('test2', bbox_inches='tight')
    plt.show()
    plt.close(fig)
